import numpy as np
    

def rem(a, b):
    """
    Returns :the remainder after division. If b <> 0, then rem(a,b) = a - fix(a/b)*b. If b = 0 or b = Inf or b = -Inf, then rem returns NaN.

    Args:
       Scalar or  Matrix inputs

    Returns:
        Scalar or  Matrix  divides A by B and returns the remainder after division
    """ 
    if isinstance(a, (int, float)) and isinstance(b, (int, float)):
        # scalar input
        if b == 0 or np.isinf(b):
            return np.nan
        else:
            return a - np.fix(a/b) * b
    elif (isinstance(a, (tuple,list)) and isinstance(b, (tuple,list))) or type(a) == np.ndarray and type(b) == np.ndarray:
        # vector or matrix input
        if np.any(b == 0) or np.any(np.isinf(b)):
            return np.full(a.shape, np.nan)
        else:
            return a - np.floor(a/b) * b
    else:
        raise TypeError("Input type not supported.")


def test01():
    # scalar input
    assert rem(5, 2) == 1
    assert np.isnan(rem(5, 0))
    assert np.isnan(rem(5, np.inf))


def test02():
    # vector input
    a = np.array([5, 6, 7, 8])
    b = np.array([2, 3, 0, np.inf])
    expected = np.array([np.nan, np.nan, np.nan, np.nan])
    assert np.array_equal(np.all(rem(a, b)), np.all(expected))


def test03():
    # matrix input
    a = np.array([[5, 6]])
    b = np.array([[2, 3]])
    expected = np.array([[1, 0]])
    assert np.array_equal(np.all(rem(a, b)),  np.all(expected))


test01()
test02()
test03()
